using System.Collections.Generic;
using System.Diagnostics;
using HTCollections.Job;
using HTCollections.Unsafe;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;
using Debug = UnityEngine.Debug;
using Random = UnityEngine.Random;

namespace HTCollections.Examples.KdTree2D
{
    public class KdTree2DExample : MonoBehaviour
    {
        public int Count;
        public float RandomRadius;
        public float QueryRange;
        public GameObject Prefab;
        public Color IdleColor;
        public Color FocusColor;
        Dictionary<float3, MeshRenderer> mSkins;
        UnsafeList<float3> mPoints;
        UnsafeKdTree2D mKdTree1;
        Custom.KdTree2D mKdTree2;

        void Start()
        {
            GameObject root = new GameObject("Root");
            mSkins = new Dictionary<float3, MeshRenderer>(Count);
            mPoints = new UnsafeList<float3>(Count, Allocator.Persistent);
            for (int i = 0; i < Count; i++)
            {
                var skin = Instantiate(Prefab, root.transform).GetComponent<MeshRenderer>();
                skin.material.color = IdleColor;
                var pos = Random.insideUnitCircle * RandomRadius;
                float3 fPos = new float3(pos.x, pos.y, 0);
                skin.transform.position = fPos;
                mSkins.Add(fPos, skin);
                mPoints.Add(fPos);
            }

            mKdTree1 = new UnsafeKdTree2D(8, Allocator.Persistent);
        }

        void OnDestroy()
        {
            mPoints.Dispose();
            mKdTree1.Dispose();
        }

        void Update()
        {
            /*
             * kd tree1 使用 BurstCompile build 性能提升10倍, query 性能提升2倍
             */
            if (Input.GetKeyDown(KeyCode.Q)) //build kd tree1
            {
                JobDataWrap<UnsafeKdTree2D> kdTreeWrap =
                    new JobDataWrap<UnsafeKdTree2D>(mKdTree1, Allocator.TempJob);
                JobDataWrap<UnsafeList<float3>> pointsWrap =
                    new JobDataWrap<UnsafeList<float3>>(mPoints, Allocator.TempJob);

                Stopwatch sw = new Stopwatch();
                sw.Start();

                KdTreeBuildJob buildJob = new KdTreeBuildJob()
                {
                    KdTree = kdTreeWrap,
                    Points = pointsWrap,
                };
                var jobHandle = buildJob.Schedule();
                jobHandle.Complete();
                mKdTree1 = kdTreeWrap.Value;

                sw.Stop();
                Debug.Log($"Kd tree build time : {sw.Elapsed.TotalMilliseconds} ms");

                kdTreeWrap.Dispose();
                mPoints = pointsWrap.Value;
                pointsWrap.Dispose();
                Debug.Log("kd tree build done");
            }

            if (Input.GetKey(KeyCode.W) && Input.GetMouseButtonDown(0)) //query kd tree1
            {
                if (Physics.Raycast(Camera.main.ScreenPointToRay(Input.mousePosition),
                        out var hit, Mathf.Infinity, 1 << 0))
                {
                    var pos = hit.point;
                    var fPos = new float3(pos.x, pos.y, 0);
                    NativeList<float3> result = new NativeList<float3>(8, Allocator.TempJob);

                    Stopwatch sw = new Stopwatch();
                    sw.Start();

                    KdTreeQueryJob queryJob = new KdTreeQueryJob()
                    {
                        KdTree = mKdTree1,
                        Pos = fPos,
                        Range = QueryRange,
                        Result = result,
                    };
                    var jobHandle = queryJob.Schedule();
                    jobHandle.Complete();

                    sw.Stop();
                    Debug.Log($"Kd tree query time : {sw.Elapsed.TotalMilliseconds} ms");

                    foreach (var item in mSkins)
                        item.Value.material.color = IdleColor;
                    foreach (var item in result)
                    {
                        if (mSkins.TryGetValue(item, out var skin))
                            skin.material.color = FocusColor;
                    }

                    result.Dispose();
                }
            }

            if (Input.GetKey(KeyCode.E) && Input.GetMouseButtonDown(0)) //query rect kd tree1
            {
                if (Physics.Raycast(Camera.main.ScreenPointToRay(Input.mousePosition),
                        out var hit, Mathf.Infinity, 1 << 0))
                {
                    var pos = hit.point;
                    var fPos = new float3(pos.x, pos.y, 0);
                    float4 rect = new float4(fPos.x - QueryRange * 0.5f, fPos.y - QueryRange * 0.5f,
                        fPos.x + QueryRange * 0.5f, fPos.y + QueryRange * 0.5f);
                    UnsafeList<float3> result = new UnsafeList<float3>(8, Allocator.TempJob);
                    JobDataWrap<UnsafeList<float3>> resultWrap =
                        new JobDataWrap<UnsafeList<float3>>(result, Allocator.TempJob);

                    Stopwatch sw = new Stopwatch();
                    sw.Start();

                    KdTreeQueryRectJob queryJob = new KdTreeQueryRectJob()
                    {
                        KdTree = mKdTree1,
                        Rect = rect,
                        Result = resultWrap,
                    };
                    var jobHandle = queryJob.Schedule();
                    jobHandle.Complete();

                    sw.Stop();
                    Debug.Log($"Kd tree query rect time : {sw.Elapsed.TotalMilliseconds} ms");

                    result = resultWrap.Value;
                    foreach (var item in mSkins)
                        item.Value.material.color = IdleColor;
                    foreach (var item in result)
                    {
                        if (mSkins.TryGetValue(item, out var skin))
                            skin.material.color = FocusColor;
                    }

                    resultWrap.Dispose();
                    result.Dispose();
                }
            }

            if (Input.GetKeyDown(KeyCode.A)) //build kd tree2
            {
                Vector3[] points = new Vector3[mPoints.Length];
                for (int i = 0; i < mPoints.Length; i++)
                    points[i] = mPoints[i];

                Stopwatch sw = new Stopwatch();
                sw.Start();

                mKdTree2 = new Custom.KdTree2D(8);
                mKdTree2.Build(points);

                sw.Stop();
                Debug.Log($"Kd tree2 build time : {sw.Elapsed.TotalMilliseconds} ms");
            }

            if (Input.GetKey(KeyCode.S) && Input.GetMouseButtonDown(0)) //query kd tree2
            {
                if (Physics.Raycast(Camera.main.ScreenPointToRay(Input.mousePosition),
                        out var hit, Mathf.Infinity, 1 << 0))
                {
                    var pos = hit.point;
                    pos = new Vector3(pos.x, pos.y, 0);
                    List<Vector3> result = new List<Vector3>(8);

                    Stopwatch sw = new Stopwatch();
                    sw.Start();

                    mKdTree2.QueryNeighborList(pos, QueryRange, ref result);

                    sw.Stop();
                    Debug.Log($"Kd tree2 query time : {sw.Elapsed.TotalMilliseconds} ms");

                    foreach (var item in mSkins)
                        item.Value.material.color = IdleColor;
                    foreach (var item in result)
                    {
                        if (mSkins.TryGetValue(item, out var skin))
                            skin.material.color = FocusColor;
                    }
                }
            }

            if (Input.GetKey(KeyCode.D) && Input.GetMouseButtonDown(0)) //query rect kd tree2
            {
                if (Physics.Raycast(Camera.main.ScreenPointToRay(Input.mousePosition),
                        out var hit, Mathf.Infinity, 1 << 0))
                {
                    var pos = hit.point;
                    pos = new Vector3(pos.x, pos.y, 0);
                    Vector4 rect = new Vector4(pos.x - QueryRange * 0.5f, pos.y - QueryRange * 0.5f,
                        pos.x + QueryRange * 0.5f, pos.y + QueryRange * 0.5f);
                    List<Vector3> result = new List<Vector3>(8);

                    Stopwatch sw = new Stopwatch();
                    sw.Start();

                    mKdTree2.QueryNeighborList(rect, ref result);

                    sw.Stop();
                    Debug.Log($"Kd tree2 query rect time : {sw.Elapsed.TotalMilliseconds} ms");

                    foreach (var item in mSkins)
                        item.Value.material.color = IdleColor;
                    foreach (var item in result)
                    {
                        if (mSkins.TryGetValue(item, out var skin))
                            skin.material.color = FocusColor;
                    }
                }
            }

            if (Input.GetKeyDown(KeyCode.Z)) //使用指针给结构体赋值：性能提升40%~50%
            {
                int loopCnt = 1000000;
                unsafe
                {
                    UnsafeTestData* datas = UnsafeUtil.Malloc<UnsafeTestData>(20, Allocator.Persistent);

                    Stopwatch sw1 = new Stopwatch();
                    sw1.Start();

                    for (int i = 0; i < loopCnt; i++)
                    {
                        datas[7] = new UnsafeTestData()
                        {
                            Index = i,
                            Value1 = 20f,
                            Value2 = 2,
                            Value3 = i,
                            Value4 = 4,
                            Value5 = i,
                            Value6 = 6,
                            Value7 = true,
                        };
                    }

                    sw1.Stop();
                    Debug.Log($"set1 : {sw1.Elapsed.TotalMilliseconds}ms");

                    Stopwatch sw2 = new Stopwatch();
                    sw2.Start();

                    for (int i = 0; i < loopCnt; i++)
                    {
                        UnsafeTestData* ptr = &datas[6];
                        ptr->Index = i;
                        ptr->Value1 = 20f;
                        ptr->Value2 = 2;
                        ptr->Value3 = i;
                        ptr->Value4 = 4;
                        ptr->Value5 = i;
                        ptr->Value6 = 6;
                        ptr->Value7 = true;
                    }

                    sw2.Stop();
                    Debug.Log($"set2 : {sw2.Elapsed.TotalMilliseconds}ms");
                }
            }
        }

        public struct UnsafeTestData
        {
            public int Index;
            public float Value1;
            public short Value2;
            public int Value3;
            public int Value4;
            public int Value5;
            public int Value6;
            public bool Value7;
        }
    }
}